import argparse
import numpy as np

import mindspore.nn as nn
from mindspore import nn, Tensor, Model, context
from mindspore import load_checkpoint, load_param_into_net
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor

from lenet import LeNet5
from resnet import resnet18
from resnet import resnet50
from animal_dataset import CreateAnimalDataset

parser = argparse.ArgumentParser(description='Image classification')
parser.add_argument('--train', action='store_true', help='run training.')
parser.add_argument('--test', action='store_true', help='run testing.')
parser.add_argument('--save_ckpt_path', type=str, default="./", help='run testing.')
parser.add_argument('--load_ckpt_path', type=str, default=None, help='run testing.')
args = parser.parse_args()


def train_feed():
    # net_with_criterion = nn.SoftmaxCrossEntropyWithLogits(reduction='mean')
    # optimizer = nn.Momentum(net.trainable_params(), learning_rate=0.01, momentum=0.9)
    # train_net = nn.TrainOneStepCell(net_with_criterion, optimizer)
    # train_net.set_train()

    # epochs = 5
    # for epoch in range(epochs):
    #     for i, data in enumerate(train_dataset.create_dict_iterator(num_epochs=epochs)):
    #         loss_val = train_net(data["image"], data["label"])
    #         if i % 10 == 0:
    #             print("Epoch {}, step {}, loss = {}".format(epoch, i, loss_val))
    #     print(loss_val)
    pass


def train(net, net_name):
    train_dataset = CreateAnimalDataset("./animals_dataset", "train")

    loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
    opt = nn.Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9)
    model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'})

    epochs = 20

    batch_num = train_dataset.get_dataset_size()
    config_ck = CheckpointConfig(save_checkpoint_steps=batch_num, keep_checkpoint_max=35)
    ckpoint_cb = ModelCheckpoint(prefix=net_name+"_train_animal", directory=args.save_ckpt_path, config=config_ck)
    loss_cb = LossMonitor(per_print_times=1)
 
    print("====================== start training ======================")
    model.train(epochs, train_dataset, callbacks=[ckpoint_cb, loss_cb])
    print("======================  end training  ======================")


def test(net):
    if not args.load_ckpt_path:
        raise RuntimeError("specify ckpt first")

    param_dict = load_checkpoint(args.load_ckpt_path)
    load_param_into_net(net, param_dict)
    test_dataset = CreateAnimalDataset("./animals_dataset", "test")

    loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
    opt = nn.Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9)
    # opt = nn.Adam(filter(lambda x: x.requires_grad, net.get_parameters()))
    model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'})

    print("====================== start testing ======================")
    res = model.eval(test_dataset)
    print("======================  end testing  ======================")
    print("result: ", res)


if __name__ == "__main__":
    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
    context.set_context(device_id=0)

    net, net_name = LeNet5(num_class=3, num_channel=3), "lenet"
    #net, net_name = resnet18(class_num=3), "resnet18"
    #net, net_name = resnet50(class_num=3), "resnet50"

    if args.train:
        train(net, net_name)
    if args.test:
        test(net)

    '''
    # 执行训练
    python train_model.py --train --save_ckpt_path=保存ckpt的地方

    # 执行推理
    python train_model.py --test --load_ckpt_path=ckpt文件路径
    '''
